package crypto

import (
	"bytes"
	"crypto/aes"
	"encoding/base64"
	"errors"
)

/**
 * AesEcbEncrypt，AES 的 ECB 模式加密
 * Params：
 *   plainText：待加密字符串的字节切片
 *   key：加密秘钥，长度必须为 16 or 24 or 32
 * Return：
 *   加密结果字节切片
 */
func AesEcbEncrypt(plainText, key []byte) ([]byte, error) {
	// 创建初始块
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	// 添加补全码
	size := block.BlockSize()
	padding := size - len(plainText)%size
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	plainText = append(plainText, padText...)
	// 将明文字符串按照初始块的长度均分为若干段，分别进行加密组成加密切片返回
	cipherText := make([]byte, len(plainText))
	for bs, be := 0, size; bs < len(plainText); bs, be = bs+size, be+size {
		block.Encrypt(cipherText[bs:be], plainText[bs:be])
	}
	return cipherText, nil
}

/**
 * AesEcbEncrypt，获取 ECB 模式加密后的 Base64 字符串
 * Params：
 *   plainText：待加密字符串的字节切片
 *   key：加密秘钥，长度必须为 16 or 24 or 32
 * Return：
 *   加密结果进行 Base64 编码后的字符串
 */
func AesEcbEncryptBase64(plainText, key []byte) (string, error) {
	encrypted, err := AesEcbEncrypt(plainText, key)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(encrypted), nil
}

/**
 * AesEcbEncrypt，AES 的 ECB 模式解密
 * Params：
 *   cipherText：ECB 模式下加密的数据字节切片
 *   key：加密秘钥，长度必须为 16 or 24 or 32
 * Return：
 *   解密后的明文字符串的字节切片
 */
func AesEcbDecrypt(cipherText, key []byte) ([]byte, error) {
	// 创建初始块
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	// 按照初始块的长度将加密切片分成若干段分别进行解密，组成解密切片
	decrypted := make([]byte, len(cipherText))
	size := block.BlockSize()
	for bs, be := 0, size; bs < len(cipherText); bs, be = bs+size, be+size {
		block.Decrypt(decrypted[bs:be], cipherText[bs:be])
	}
	// 去除补全码
	length := len(decrypted)
	unPadding := int(decrypted[length-1])
	if length-unPadding < 0 {
		return nil, errors.New("encryption error")
	}
	return decrypted[:(length - unPadding)], nil
}

/**
 * AesEcbDecryptBase64，根据 ECB 模式加密结果进行 base64 编码后的字符串获取解密明文字符串的字节切片
 * Params：
 *   cipherTextBase64：ECB 模式加密结果进行 base64 编码后的字符串
 *   key：加密秘钥，长度必须为 16 or 24 or 32
 * Return：
 *   解密后的明文字符串字节切片
 */
func AesEcbDecryptBase64(cipherTextBase64 string, key []byte) ([]byte, error) {
	cipherText, err := base64.StdEncoding.DecodeString(cipherTextBase64)
	if err != nil {
		return nil, err
	}
	return AesEcbDecrypt(cipherText, key)
}
